import json
import os
import pandas as pd
from tqdm import tqdm

from model_api import qwen3_reranker_api


path_dir = os.path.dirname(os.path.abspath(__file__))

def call_rerank_main(qas):
    question, answer = qas
    res = qwen3_reranker_api.call_api(question, answer)
    return res[0]['relevance_score']

def build_prop_to_desc(props_with_desc):
    prop2desc = dict()
    for prop_with_desc in props_with_desc:
        prop2desc[prop_with_desc['prop']] = prop_with_desc['desc']
    return prop2desc

if __name__ == '__main__':

    phase25_model = 'qwen3-max'
    with open(os.path.join(path_dir, "phase2.5_output", f"{phase25_model}_prop_desc.jsonl"), 'r', encoding='utf-8') as fp:
        lines = fp.readlines()
        props_with_desc = [json.loads(line) for line in lines]
    prop2desc = build_prop_to_desc(props_with_desc)
    phase4_model = 'qwen3-max'
    with open(os.path.join(path_dir, "phase4_output", f"{phase4_model}_questions_hard.json"), "r", encoding='utf-8') as fp:
        hard_qas = json.load(fp)
    for hard_qa in tqdm(hard_qas):
        score = call_rerank_main((hard_qa['question'], hard_qa['answer']))
        hard_qa['score'] = score
    for hard_qa in tqdm(hard_qas):
        question = hard_qa['question'].rstrip() + '\n'
        cnt = 1
        for prop in hard_qa['prop']:
            if prop in prop2desc:
                question += f"参考资料{cnt}：{prop2desc[prop]}\n"
                cnt += 1
        score_with_ref = call_rerank_main((question, hard_qa['answer']))
        hard_qa['score_with_ref'] = score_with_ref
    output_fp = open(os.path.join(path_dir, "phase9_output", f"{phase4_model}_questions_hard.json"), 'w', encoding='utf-8')
    json.dump(hard_qas, output_fp, ensure_ascii=False, indent=4)
    output_fp.close()

    with open(os.path.join(path_dir, "phase9_output", f"qwen3-max_questions_hard.json"), 'r', encoding='utf-8') as fp:
        hard_qas = json.load(fp)
    for hard_qa in tqdm(hard_qas):
        props = hard_qa.pop('prop')
        hard_qa['prop1'] = props[0]
        hard_qa['prop2'] = props[1]
    df = pd.DataFrame(hard_qas)
    df.to_excel(os.path.join(path_dir, "phase9_output", f"qwen3-max_questions_hard.xlsx"), index=False)

